# -*- coding:utf-8 -*-
"""
双色球(SSQ)数据专属爬取脚本
支持通过--data-type参数指定爬取数据用途(训练/预测), 分别保存不同文件
预测模式仅保留最近89期数据, 训练模式保留全量数据
纯文本日志输出, 无图标
"""
import os
import time
import csv
import requests
import argparse  # 新增: 命令行参数解析
from bs4 import BeautifulSoup
import traceback
# 从config导入全局日志和配置
from config.ssq_config import init_global_logger, logger, BASE_URL, TIMEOUT, RETRY_TIMES, DELAY, SSQ_CONFIG


class SSQLotterySpider:
    def __init__(self, run_type: str):
        """初始化爬虫: 加载配置、创建目录、根据数据类型确定保存路径"""
        # 校验数据类型合法性
        if run_type not in ["train", "predict"]:
            raise ValueError(f"数据类型错误！仅支持 'train' 或 'predict', 当前输入: {run_type}")
        
        self.run_type = run_type  # 新增: 记录数据用途(train/predict)
        self.predict_recent_num = 89  # 新增: 预测模式保留最近89期数据
        self.full_data_url = BASE_URL + SSQ_CONFIG["url_pattern"]
        
        # 新增: 根据数据类型确定保存文件名(从配置文件读取, 避免硬编码)
        self.data_folder = SSQ_CONFIG["train_file"]["DATA_FOLDER"]
        self.csv_path = os.path.join(self.data_folder,SSQ_CONFIG["train_file"]["csv_file"])
        if self.run_type == "predict":
            self.data_folder = SSQ_CONFIG["predict_file"]["DATA_FOLDER"]
            self.csv_path = os.path.join(self.data_folder,SSQ_CONFIG["predict_file"]["csv_file"])

        # 创建数据存储目录(不存在则自动创建)
        if not os.path.exists(self.data_folder):
            os.makedirs(self.data_folder)
            logger.info(f"创建数据存储目录: {os.path.abspath(self.data_folder)}")
        
        init_global_logger(log_file=f"spider_{run_type}_data.log", pid=os.getpid())  # 按用途区分日志文件

        # 初始化日志信息
        logger.info(f"双色球爬虫初始化完成(数据用途: {run_type})")
        logger.debug(f"全量数据请求URL: {self.full_data_url}")
        logger.debug(f"数据保存路径: {self.csv_path}")
        if self.run_type == "predict":
            logger.info(f"预测模式: 仅保留最近{self.predict_recent_num}期数据")

    def _send_request(self) -> str:
        """发送请求获取全量数据页面(逻辑不变)"""
        logger.info(f"开始请求全量数据URL: {self.full_data_url}")
        
        # 构造请求头
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
            "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
            "Referer": BASE_URL
        }

        # 重试逻辑
        for retry in range(1, RETRY_TIMES + 1):
            try:
                response = requests.get(
                    url=self.full_data_url,
                    headers=headers,
                    verify=False,
                    timeout=TIMEOUT
                )
                response.encoding = "utf-8"

                if response.status_code == 200:
                    # 校验页面有效性
                    if len(response.text) < 1024 * 100:
                        logger.warning(f"警告: 请求返回页面过小({len(response.text)/1024:.1f}KB), 可能是反爬页面, 重试第{retry}次")
                        time.sleep(DELAY * retry)
                        continue
                    logger.success(f"{SSQ_CONFIG['name']}请求成功, 页面大小: {len(response.text)/1024:.1f}KB")
                    time.sleep(DELAY)
                    return response.text
                else:
                    logger.warning(f"警告: {SSQ_CONFIG['name']}请求失败, 状态码: {response.status_code}, 重试第{retry}次")

            except requests.exceptions.Timeout:
                logger.error(f"错误: {SSQ_CONFIG['name']}请求超时({TIMEOUT}秒), 重试第{retry}次")
            except requests.exceptions.ConnectionError:
                logger.error(f"错误: {SSQ_CONFIG['name']}网络连接异常, 重试第{retry}次")
            except Exception as e:
                logger.error(f"错误: {SSQ_CONFIG['name']}请求未知异常: {str(e)},\n{traceback.format_exc()}, 重试第{retry}次")

            time.sleep(DELAY * retry)

        # 多次重试失败后抛出异常
        raise Exception(
            f"{SSQ_CONFIG['name']}请求失败(已重试{RETRY_TIMES}次), 请按以下步骤排查: \n"
            f"1. 手动访问URL: {self.full_data_url}, 确认页面是否显示数据；\n"
            f"2. 切换网络(如手机热点), 排查IP是否被封禁；\n"
            f"3. 检查config.py中的User-Agent是否为当前浏览器真实UA。"
        )

    def _parse_page(self, html: str) -> dict:
        """解析全量数据页面, 提取双色球历史数据(逻辑不变)"""
        logger.info("开始解析全量数据页面...")
        parsed_data = {}
        soup = BeautifulSoup(html, "lxml")

        # 定位数据表格
        tbody = soup.find("tbody", attrs={"id": SSQ_CONFIG["parser_rules"]["tbody_id"]})
        if not tbody:
            logger.error(f"错误: 未找到数据表格(tbody#{SSQ_CONFIG['parser_rules']['tbody_id']}), 页面结构可能变化！")
            # 输出页面中所有tbody的ID, 帮助用户修正配置
            all_tbody = soup.find_all("tbody")
            if all_tbody:
                tbody_ids = [tb.get("id") for tb in all_tbody if tb.get("id")]
                logger.error(f"页面中已存在的tbody ID: {tbody_ids}(请核对config.py中的parser_rules)")
            return parsed_data

        # 遍历表格行, 提取数据
        tr_list = tbody.find_all("tr")
        if not tr_list:
            logger.warning("警告: 数据表格中无任何行数据, 可能是反爬页面或无历史数据")
            return parsed_data
        logger.debug(f"页面共包含{len(tr_list)}期{SSQ_CONFIG['name']}数据")

        for tr in tr_list:
            try:
                td_list = tr.find_all("td")
                # 校验行数据完整性
                min_td_count = SSQ_CONFIG["parser_rules"]["date_col_idx"] + 1
                if len(td_list) < min_td_count:
                    logger.warning(f"警告: 跳过残缺行(TD数量不足{min_td_count}个): {tr.text.strip()}")
                    continue

                # 提取核心数据
                issue = int(td_list[SSQ_CONFIG["parser_rules"]["issue_col_idx"]].text.strip())  # 期号
                date = td_list[SSQ_CONFIG["parser_rules"]["date_col_idx"]].text.strip()        # 日期
                red_balls = [int(td_list[idx].text.strip()) for idx in SSQ_CONFIG["parser_rules"]["red_ball_col_idxs"]]  # 红球
                blue_ball = int(td_list[SSQ_CONFIG["parser_rules"]["blue_ball_col_idx"]].text.strip())  # 蓝球

                # 组装数据
                parsed_data[issue] = [date] + red_balls + [blue_ball]
                logger.debug(f"解析成功: 期号{issue} → {parsed_data[issue]}")
                # 每解析100期打印进度
                if len(parsed_data) % 100 == 0:
                    logger.debug(f"已解析{len(parsed_data)}期数据, 最新期号: {issue}")

            except ValueError as e:
                logger.error(f"错误: 数据格式错误(非整数): {str(e)},\n{traceback.format_exc()}, 行内容: {tr.text.strip()}")
            except IndexError as e:
                logger.error(f"错误: 列索引越界: {str(e)},\n{traceback.format_exc()}, 请核对config.py中的parser_rules, 行内容: {tr.text.strip()}")
            except Exception as e:
                logger.error(f"错误: 行解析未知异常: {str(e)},\n{traceback.format_exc()}, 堆栈信息: {traceback.format_exc()}")

        # 解析完成, 输出统计信息
        if parsed_data:
            sorted_issues = sorted(parsed_data.keys())
            logger.success(
                f"页面解析完成！共获取{len(parsed_data)}期{SSQ_CONFIG['name']}数据\n"
                f"时间范围: 第{sorted_issues[0]}期({parsed_data[sorted_issues[0]][0]})~ 第{sorted_issues[-1]}期({parsed_data[sorted_issues[-1]][0]})"
            )
        return parsed_data

    def _save_data(self, parsed_data: dict) -> None:
        """保存数据到CSV(预测模式仅保留最近89期)"""
        if not parsed_data:
            raise ValueError("错误: 无有效数据可保存, 解析过程未获取到双色球数据")

        # 1. 对原始期号按升序排序(保证时序正确)
        sorted_original_issues = sorted(parsed_data.keys())
        
        # 新增: 预测模式筛选最近89期数据
        if self.run_type == "predict":
            total_parsed = len(sorted_original_issues)
            if total_parsed < self.predict_recent_num:
                logger.warning(f"警告: 解析到的总期数({total_parsed})不足{self.predict_recent_num}期, 将保留全部数据")
            else:
                # 截取最近89期(排序后取后89个)
                sorted_original_issues = sorted_original_issues[-self.predict_recent_num:]
                logger.info(f"预测模式: 筛选最近{self.predict_recent_num}期数据, 期号范围: {sorted_original_issues[0]}~{sorted_original_issues[-1]}")

        # 2. 生成新的顺序期号(1,2,3...), 与原始期号一一对应
        new_issue_mapping = {original_issue: i+1 for i, original_issue in enumerate(sorted_original_issues)}

        # 先删除同类型旧文件(避免数据覆盖冲突)
        if os.path.exists(self.csv_path):
            os.remove(self.csv_path)
            logger.debug(f"删除旧{self.run_type}数据文件: {self.csv_path}")

        try:
            with open(self.csv_path, "w", newline="", encoding="utf-8") as f:
                writer = csv.writer(f)
                # 写入表头(保持不变, idx现在代表顺序编号)
                writer.writerow(SSQ_CONFIG["csv_header"])
                # 3. 按新期号顺序写入数据(原始期号已排序, 新期号自然递增)
                for original_issue in sorted_original_issues:
                    new_issue = new_issue_mapping[original_issue]  # 获取顺序编号
                    # 写入数据: 新期号 + 原始数据(日期+红球+蓝球)
                    writer.writerow([new_issue] + parsed_data[original_issue])
            logger.success(
                f"{self.run_type}数据保存成功: {os.path.abspath(self.csv_path)} "
                f"(共{len(sorted_original_issues)}期, 新期号从1到{len(sorted_original_issues)})"
            )
        except Exception as e:
            raise Exception(f"错误: {self.run_type}数据保存失败: {str(e)},\n{traceback.format_exc()}, 堆栈信息: {traceback.format_exc()}")

    def run(self) -> None:
        """爬虫主流程: 请求→解析→保存(逻辑不变)"""
        logger.info(f"双色球{self.run_type}数据爬取任务启动")
        full_data = {}

        try:
            # 1. 请求全量数据页面
            html = self._send_request()
            # 2. 解析页面数据
            full_data = self._parse_page(html)
            # 3. 保存数据
            if full_data:
                self._save_data(full_data)
                logger.success(f"双色球{self.run_type}数据爬取任务完成！")
                return True
            logger.warning(f"警告: {self.run_type}数据爬取任务完成, 但未获取到有效数据(原因见解析日志)")
            return False

        except Exception as e:
            logger.error(f"错误: {self.run_type}数据爬虫任务异常终止: {str(e)},\n{traceback.format_exc()}")
            # 异常时保存已解析的部分数据
            raise Exception(f"错误: {self.run_type}数据爬虫任务异常终止: {str(e)},\n{traceback.format_exc()}")


def main():
    """新增: 命令行参数解析+脚本入口"""
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="双色球数据爬取脚本(支持训练/预测数据区分)")
    parser.add_argument(
        "--data-type",
        required=True,
        choices=["train", "predict"],
        help="数据用途: train(训练数据, 保留全量)、predict(预测数据, 仅保留最近89期)"
    )
    args = parser.parse_args()

    # 启动爬虫
    try:
        ssq_spider = SSQLotterySpider(run_type=args.run_type)
        ssq_spider.run()
    except Exception as e:
        logger.critical(f"错误: 脚本运行失败: {str(e)},\n{traceback.format_exc()}")
        exit(1)


if __name__ == "__main__":
    main()